///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/onepdm.cpp
///
/// \brief   implementation file for the monte carlo one particle density matrix
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<vector>
#include<string>
#include<numeric>
#include<cassert>
#include<algorithm>

#include<boost/format.hpp>
#include<boost/shared_ptr.hpp>

#include<formic/all_data.h>
#include<formic/fqmc/onepdm.h>
#include<formic/wfn/configuration.h>
#include<formic/fqmc/vmc_walker.h>
#include<formic/fqmc/sampler.h>
#include<formic/wfn/wfn.h>
#include<formic/input/base.h>
#include<formic/mpi/interface.h>
#include<formic/numeric/numeric.h>
#include<formic/output/output.h>
#include<formic/exception.h>
#include<formic/lapack/interface.h>
#include<formic/random/random.h>

namespace formic { namespace fqmc {

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  a class to accumulate the one particle RDM during a markov chain evolution
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////
  template <class S> class OnepdmAcc : public formic::Accumulator<S> {

    // data members
    private:

      /// \brief  number of atomic orbitals
      int _nao;

//      /// \brief  pointer to the guiding function
//      boost::shared_ptr<formic::CoeffObj<S> > _gf;

      /// \brief  reference to the walker's configuration
      const formic::Configuration & _config;

      /// \brief  the alpha 1-particle rdm
      std::vector<S> _ardm;

      /// \brief  the beta 1-particle rdm
      std::vector<S> _brdm;

      /// \brief  the alpha natural orbitals
      std::vector<S> _ano;

      /// \brief  the beta natural orbitals
      std::vector<S> _bno;

      /// \brief  the alpha natural orbital occupations
      std::vector<double> _anocc;

      /// \brief  the beta natural orbital occupations
      std::vector<double> _bnocc;

    // private member functions
    private:

    public:

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  constructs the accumulator around a specific coefficient object
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      OnepdmAcc(const formic::InputBase & userinp,
                const boost::shared_ptr<formic::CoeffObj<S> > & gf,
                const boost::shared_ptr<formic::CoeffObj<S> > & vf,
                const formic::Configuration & config)
        : formic::Accumulator<S>(gf, vf),
          _nao(userinp.get<int>("nao")),
          _config(config),
          _ardm(userinp.get<int>("nao") * userinp.get<int>("nao"), formic::zero(S())),
          _brdm(userinp.get<int>("nao") * userinp.get<int>("nao"), formic::zero(S())),
          _ano(userinp.get<int>("nao") * userinp.get<int>("nao"), formic::zero(S())),
          _bno(userinp.get<int>("nao") * userinp.get<int>("nao"), formic::zero(S())),
          _anocc(userinp.get<int>("nao"), 0.0),
          _bnocc(userinp.get<int>("nao"), 0.0)
      {}

      /// \brief  return the accumulator's name
      std::string name() const { return "one-pdm accumulator"; }

      /// \brief  returns a reference to the alpha natural orbitals
      const std::vector<S> & ano() const { return _ano; }

      /// \brief  returns a reference to the beta natural orbitals
      const std::vector<S> & bno() const { return _bno; }

      /// \brief  returns a reference to the alpha natural orbital occupations
      const std::vector<double> & anocc() const { return _anocc; }

      /// \brief  returns a reference to the beta natural orbital occupations
      const std::vector<double> & bnocc() const { return _bnocc; }

//      ///////////////////////////////////////////////////////////////////////////////////////////////////
//      /// \brief  Prints some statistics about the most recent sample.
//      ///
//      ///////////////////////////////////////////////////////////////////////////////////////////////////
//      void print_statistics(formic::Output<std::ofstream> & fout) {
//
//        if ( formic::mpi::rank() == 0 ) {
//          fout << boost::format("gradient accumulation statistics on process %i:") % formic::mpi::rank() << std::endl;
//          fout << std::endl;
//          fout << boost::format("   acceptance ratio = %20.12f") % (this->_naccept / this->_total_weight) << std::endl;
//          fout << boost::format("      gradient norm = %20.12f") % _grad_norm << std::endl;
//          fout << boost::format("             energy = %20.12f") % _energy << std::endl;
//          fout << boost::format("    energy variance = %20.12f") % _variance << std::endl;
//          fout << boost::format("   energy std. dev. = %20.12f") % std::sqrt(_variance) << std::endl;
//          fout << std::endl;
//        }
//
//      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief Accumulate the local rdm contributions
      ///
      /// \param[in]      userinp  the user's input options
      /// \param[in]      w        the weight to apply to this configuration's contribution
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void child_accumulate(const formic::InputBase & userinp, const double w) {

        // get an object to store configuration differences
        formic::ConfigDiff cdiff(userinp);
        cdiff.newly_full().assign(1,0);
        cdiff.newly_empty().assign(1,0);

        // get rdm contributions
        for (int b = 0; b < 2; b++) {

          // get reference to rdm
          std::vector<S> & rdm = ( b == 0 ? _ardm : _brdm );

          // get references to occupied and unoccupied orbital lists
          const std::vector<int> & occ_to_site   = ( b == 0 ? _config.a_occ_to_site()     : _config.b_occ_to_site()   );
          const std::vector<int> & unocc_to_site = ( b == 0 ? _config.a_unocc_to_site()   : _config.b_unocc_to_site() );

          // get number of particles and holes
          const int np = occ_to_site.size();
          const int nu = unocc_to_site.size();

          // loop over occupieds
          for (std::vector<int>::const_iterator oon = occ_to_site.begin(); oon != occ_to_site.end(); oon++) {

            // diagonal term
            rdm.at((*oon)+(*oon)*_nao) += w * this->_vgs * formic::unity(S());

            // set the orbital to be vacated
            cdiff.newly_empty().at(0) = 2 * (*oon) + b;

            // off-diagonal terms (loop over unoccupieds)
            for (std::vector<int>::const_iterator uon = unocc_to_site.begin(); uon != unocc_to_site.end(); uon++) {

              // set the orbital to be filled
              cdiff.newly_full().at(0) = 2 * (*uon) + b;

              // get contribution to the rdm element
              rdm.at((*oon)+(*uon)*_nao) += w * this->_vgs * this->_base_vf->get_ratio(cdiff);

            }

          }

        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief finalizes the accumulation (this may include for example communicating the results)
      ///
      /// \param[in]      userinp  the user's input options
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void child_finalize(const formic::InputBase & userinp) {

        // get some more work space
        int lwork = _nao*_nao + 3*_nao;
        std::vector<S> work(lwork, formic::zero(S()));
        std::vector<double> rwork(lwork, 0.0);

        // do both alpha and beta RDMs
        for (int b = 0; b < 2; b++) {

          // get reference to rdm, natural orbital array, and natural orbital occupations
          std::vector<S> & rdm  = ( b == 0 ? _ardm  : _brdm  );
          std::vector<S> & no   = ( b == 0 ? _ano   : _bno   );
          std::vector<double> & nocc = ( b == 0 ? _anocc : _bnocc );

          // divide the RDM by the total weight
          formic::xscal(rdm.size(), 1.0 / this->_total_weight / this->_vgsa, &rdm[0], 1);

          // reduce the rdm
          formic::mpi::allreduce(&rdm[0], &work[0], rdm.size(), MPI::SUM);
          formic::xcopy(rdm.size(), &work[0], 1, &rdm[0], 1);

          // symmetrize the rdm
          for (int i = 0; i < _nao; i++)
          for (int j = 0; j < i;    j++) {
            const S val = 0.5 * ( rdm.at(i+j*_nao) + formic::conj(rdm.at(j+i*_nao)) );
            rdm.at(i+j*_nao) = val;
            rdm.at(j+i*_nao) = formic::conj(val);
          }

          // diagonalize the rdm
          int info = 0;
          formic::xcopy(rdm.size(), &rdm[0], 1, &no[0], 1);
          formic::xsyev('V', 'U', _nao, &no[0], _nao, &nocc[0], &work[0], 5*_nao, &rwork[0], info);
          if (info != 0 )
            throw formic::Exception("1-body rdm diagonalization (b=%i) routine xsyev failed with error code %i") % b % info;

          // reverse the order of the natural orbitals
          for (int i = 0; i < _nao/2; i++) {
            const int j = _nao - i - 1;
            std::swap(nocc[i], nocc[j]);
            for (int k = 0; k < _nao; k++)
              std::swap(no[k+i*_nao], no[k+j*_nao]);
          }

          // transpose the natural orbital coefficient matrix so that the natural orbitals are in the rows
          for (int i = 0; i < _nao; i++)
          for (int j = 0; j < i; j++)
            std::swap(no[i+j*_nao], no[j+i*_nao]);

        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  resets the accumulator
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void child_reset(const formic::Configuration & config) {
        std::fill(_ardm.begin(), _ardm.end(), formic::zero(S()));
        std::fill(_brdm.begin(), _brdm.end(), formic::zero(S()));
        std::fill(_ano.begin(), _ano.end(), formic::zero(S()));
        std::fill(_bno.begin(), _bno.end(), formic::zero(S()));
        std::fill(_anocc.begin(), _anocc.end(), 0.0);
        std::fill(_bnocc.begin(), _bnocc.end(), 0.0);
      }

  };

} // end namespace fqmc

} // end namespace formic

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  get the coefficient object's natural orbitals
///
/// \param[in]      userinp        the user's input options
/// \param[in]      burn_len       the burn-in length
/// \param[in]      samp_len       the sample length
/// \param[in]      init_config    the initial configuration
/// \param[in]      rng            the random number generator to use
/// \param[in]      gf             a shared pointer to the guiding function
/// \param[in]      vf             a shared pointer to the value function
/// \param[out]     anocc          on exit, the alpha natural orbital occupations
/// \param[out]     bnocc          on exit, the beta natural orbital occupations
/// \param[out]     ano            on exit, the alpha natural orbitals
/// \param[out]     bno            on exit, the beta natural orbitals
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::fqmc::get_nos(const formic::InputBase & userinp,
                                             boost::shared_ptr<formic::fqmc::Sampler<S> > & sampler,
                                             const boost::shared_ptr<formic::CoeffObj<S> > & vf,
                                             std::vector<double> & anocc,
                                             std::vector<double> & bnocc,
                                             std::vector<S> & ano,
                                             std::vector<S> & bno) {

  // create an accumulator for the one body rdm
  boost::shared_ptr<formic::Accumulator<S> > racc( new formic::fqmc::OnepdmAcc<S>(userinp, sampler->gf(), vf, sampler->config()) );

  // sample the one body rdm
  sampler->simple_sample() << racc << vf;

  // get a derived class pointer to the rdm accumulator
  formic::fqmc::OnepdmAcc<S> * racc_ptr = dynamic_cast<formic::fqmc::OnepdmAcc<S> *>(racc.get());

  // extract the natural orbital occupations
  anocc = racc_ptr->anocc();
  bnocc = racc_ptr->bnocc();

  // extract the natural orbitals
  ano = racc_ptr->ano();
  bno = racc_ptr->bno();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  prints the supplied natural orbitals and occupations to the output file
///
/// \param[in]      n              the number of orbitals
/// \param[in]      name           name of the orbitals (usually alpha or beta)
/// \param[in]      occs           the natural orbital occupations
/// \param[in]      orbs           the natural orbitals
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::fqmc::nat_orb_text_print(const int n, const std::string & name, const std::vector<double> & occs, const std::vector<S> & orbs) {

  formic::of << boost::format("printing %s natural orbitals and occupations:") % name << std::endl << std::endl;
  for (int i = 0; i < n; i++)
    formic::of << " " << formic::format_number("%14.8f", occs.at(i));
  formic::of << std::endl;
  for (int i = 0; i < n; i++)
    formic::of << "---------------";
  formic::of << std::endl;
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++)
      formic::of << " " << formic::format_number("%14.8f", orbs.at(j+n*i));
    formic::of << std::endl;
  }
  formic::of << std::endl;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  compute and print the natural orbitals
///
/// \param[in]      userinp        the user's input options
/// \param[in]      burn_len       the burn-in length
/// \param[in]      samp_len       the sample length
/// \param[in]      init_config    the initial configuration
/// \param[in]      rng            the random number generator to use
/// \param[in]      gf             a shared pointer to the guiding function
/// \param[in]      vf             a shared pointer to the value function
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::fqmc::print_nos(const formic::InputBase & userinp,
                                               boost::shared_ptr<formic::fqmc::Sampler<S> > & sampler,
                                               const boost::shared_ptr<formic::CoeffObj<S> > & vf) {

  // get the natural orbitals and their occupations
  std::vector<double> aocc, bocc;
  std::vector<S> ano, bno;
  formic::fqmc::get_nos(userinp, sampler, vf, aocc, bocc, ano, bno);

  // print the natural orbitals
  if ( formic::mpi::rank() == 0 ) {
    formic::fqmc::nat_orb_text_print(userinp.get<int>("nao"), "alpha", aocc, ano);
    formic::fqmc::nat_orb_text_print(userinp.get<int>("nao"),  "beta", bocc, bno);
  }

}

template void formic::fqmc::print_nos(const formic::InputBase & userinp,
                                      boost::shared_ptr<formic::fqmc::Sampler<double> > & sampler,
                                      const boost::shared_ptr<formic::CoeffObj<double> > & vf);
template void formic::fqmc::print_nos(const formic::InputBase & userinp,
                                      boost::shared_ptr<formic::fqmc::Sampler<std::complex<double> > > & sampler,
                                      const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > & vf);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "PrintNO" command, which computes and prints a coefficient
///          object's natural orbitals.
///
///          1st argument is the name of the coefficient object to use
///          2nd argument is the burn in length
///          3rd argument is the sample length
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void formic::fqmc::PrintNO_command_function(formic::Archive & arch, formic::AllData & data) {

  // read the name of the coefficient object to use
  std::string co_name;
  formic::mpi::read_and_bcast(arch, co_name, "Failed to read coefficient object name in formic::fqmc::PrintNO_command_function.");

  // read the burn in length
  int burn_len;
  formic::mpi::read_and_bcast(arch, burn_len, "Failed to read burn in length in formic::fqmc::PrintNO_command_function.");

  // read the sample length
  int samp_len;
  formic::mpi::read_and_bcast(arch, samp_len, "Failed to read sample length in formic::fqmc::PrintNO_command_function.");

  // compute and print the natural orbitals
  if (data.userinp.get<std::string>("scalar_type") == "real") {
    typedef formic::fqmc::Sampler<double> Smplr;
    boost::shared_ptr<Smplr> sampler( new Smplr(data.userinp, burn_len, samp_len, formic::Configuration(data.userinp), formic::LaggedFibonacci<>(), data.wd.get_co(co_name)) );
    formic::fqmc::print_nos(data.userinp, sampler, data.wd.get_co(co_name));
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    typedef formic::fqmc::Sampler<std::complex<double> > Smplr;
    boost::shared_ptr<Smplr> sampler( new Smplr(data.userinp, burn_len, samp_len, formic::Configuration(data.userinp), formic::LaggedFibonacci<>(), data.wc.get_co(co_name)) );
    formic::fqmc::print_nos(data.userinp, sampler, data.wc.get_co(co_name));
  } else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"PrintNO\"") % data.userinp.get<std::string>("scalar_type");

}
